package com.gitee.wsl.ext.coroutines


import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.onClosed
import kotlinx.coroutines.channels.onFailure
import kotlinx.coroutines.channels.onSuccess
import kotlinx.coroutines.selects.onTimeout
import kotlinx.coroutines.selects.whileSelect

/**
 * Receive a batch of items from a channel, up to time and batch size limits.
 *
 * @param channel the channel to receive items from
 * @param maxTimeMillis maximum time in milliseconds to wait for items
 * @param maxSize maximum batch size
 *
 * @return the list of items in the order received
 */
@OptIn(ExperimentalCoroutinesApi::class)
 suspend fun <E> ReceiveChannel<E>.receiveBatch(
    maxTimeMillis: Long,
    maxSize: Int,
): List<E> {
    val batch = mutableListOf<E>()
    whileSelect {
        onTimeout(maxTimeMillis) { false }
        onReceiveCatching { result ->
            result
                .onFailure { if (it != null) throw it }
                .onClosed { return@onReceiveCatching false }
                .onSuccess { batch += it }
            batch.size < maxSize
        }
    }
    return batch
}
